"""
Simple class to calculate precision, recall, and NIAP
"""

class Evaluator:
    def __init__(self, result_file_name, judgment_file_name):
        self.result_file_name = result_file_name
        self.judgment_file_name = judgment_file_name
        self.retrieve = []
        self.judgment = []
        self.isrel = {}
        self.nquery = 0
        self.numrel = {}
        
        f = open(self.result_file_name, 'r')
        for line in f.readlines(): 
            qid = line.split()[0]
            did = line.split()[1]
            self.retrieve.append((int(qid), int(did)))
            if int(qid) >= self.nquery:
                self.nquery = int(qid)
        
        #print "RETRIEVE = ",self.retrieve
        
        f = open(self.judgment_file_name, 'r')
        for line in f.readlines():
            qid = line.split()[0]
            did = line.split()[1]
            self.judgment.append((int(qid), int(did)))
            self.isrel[(int(qid), int(did))] = True
            if int(qid) in self.numrel.keys():
                self.numrel[int(qid)] += 1
            else:
                self.numrel[int(qid)] = 1
                
        #print "NUMREL = ",self.numrel
            
    def calculate(self):
        ret = []
        #print "NQUERY =",self.nquery
        #print "NUMREL =",self.numrel
        #print "ISREL =",self.isrel
        for i in range(1, self.nquery + 1):
            niap, prec, recall, nrel, nret = 0.0,0.0,0.0,0,0
            j = 0
            #print "j =",j
            #print "len(self.retrieve)",len(self.retrieve)
            #print "self.retrieve[j][0]",self.retrieve[j][0]
            #print "i",i
            
            while (j < len(self.retrieve)) and (self.retrieve[j][0] <= i):
                if self.retrieve[j][0] == i:
                    if (i,self.retrieve[j][1]) in self.isrel and self.isrel[(i,self.retrieve[j][1])] == True:
                        nrel += 1
                    nret += 1
                    if (i,self.retrieve[j][1]) in self.isrel and self.isrel[(i,self.retrieve[j][1])] == True:
                        niap += nrel/float(nret)
                j += 1
                
            #print "NRET =", nret
            #print "NREL =", nrel
            #print "NUMREL[i] = ", self.numrel[i]
            prec = nrel/float(nret)
            recall = nrel/float(self.numrel[i])
            niap /= float(self.numrel[i])
            ret.append((i, prec, recall, niap))
        
        return ret

    def writeResult(self, file_name):
        """(no_query, prec, recall, NIAP)"""
        f = open(file_name, 'w')
        ret = self.calculate()
        f.writelines("query prec recall niap\n")
        for line in ret:
            f.writelines(str(line[0]) + " " + str(line[1]) + " " + str(line[2]) + " " + str(line[3]) + "\n")
            
            
if __name__ == '__main__':
    ev = Evaluator('result.txt', 'rel.txt')
    ev.writeResult('evaluation.txt')
    
        
